package com.example;

import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.AssignerWithPeriodicWatermarks;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import javax.annotation.Nullable;
import java.sql.Timestamp;
import java.util.*;

public class UserRepurchasePrediction {
    public static void main(String[] args) throws Exception {
        // 设置执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);
        env.setParallelism(1);

        DataStream<String> inputStream = env.readTextFile("data/user_purchases.csv");

        DataStream<UserPurchase> purchaseStream = inputStream
                .filter(line -> !line.startsWith("userId"))
                .map(new MapFunction<String, UserPurchase>() {
                    @Override
                    public UserPurchase map(String line) throws Exception {
                        String[] fields = line.split(",");
                        return new UserPurchase(
                                fields[0], // userId
                                fields[1], // orderId
                                Double.parseDouble(fields[3]), // amount
                                Timestamp.valueOf(fields[2])  // purchaseTime
                        );
                    }
                })
                .assignTimestampsAndWatermarks(new PurchaseTimeExtractor());

        // 按用户ID分组，计算复购指标
        DataStream<RepurchaseIndicator> indicatorStream = purchaseStream
                .keyBy(UserPurchase::getUserId)
                .timeWindow(Time.days(30))
                .aggregate(new PurchaseStatisticsAggregator(), new RepurchaseIndicatorGenerator());

        // 输出结果
        indicatorStream.print();

        // 执行作业
        env.execute("User Repurchase Prediction");
    }

    // 提取时间戳和生成水位线
    public static class PurchaseTimeExtractor implements AssignerWithPeriodicWatermarks<UserPurchase> {
        private final long maxOutOfOrderness = 3000; // 3秒
        private long currentMaxTimestamp;

        @Nullable
        @Override
        public Watermark getCurrentWatermark() {
            return new Watermark(currentMaxTimestamp - maxOutOfOrderness);
        }

        @Override
        public long extractTimestamp(UserPurchase element, long previousElementTimestamp) {
            long timestamp = element.getPurchaseTime().getTime();
            currentMaxTimestamp = Math.max(timestamp, currentMaxTimestamp);
            return timestamp;
        }
    }

    // 聚合函数，计算每个用户的购买统计信息
    public static class PurchaseStatisticsAggregator implements AggregateFunction<UserPurchase, PurchaseStatistics, PurchaseStatistics> {
        @Override
        public PurchaseStatistics createAccumulator() {
            return new PurchaseStatistics();
        }

        @Override
        public PurchaseStatistics add(UserPurchase value, PurchaseStatistics accumulator) {
            accumulator.purchaseCount++;
            accumulator.totalAmount += value.getAmount();
            accumulator.purchaseTimes.add(value.getPurchaseTime().getTime());
            return accumulator;
        }

        @Override
        public PurchaseStatistics getResult(PurchaseStatistics accumulator) {
            // 计算平均订单金额
            accumulator.averageAmount = accumulator.purchaseCount > 0 ?
                    accumulator.totalAmount / accumulator.purchaseCount : 0;

            // 计算购买频率（假设窗口大小为30天）
            if (accumulator.purchaseCount > 1) {
                Collections.sort(accumulator.purchaseTimes);
                long totalInterval = 0;
                for (int i = 1; i < accumulator.purchaseTimes.size(); i++) {
                    totalInterval += (accumulator.purchaseTimes.get(i) - accumulator.purchaseTimes.get(i-1));
                }
                // 平均间隔（天）
                accumulator.averageInterval = totalInterval / (1000.0 * 3600 * 24 * (accumulator.purchaseCount - 1));
                // 频率（次/月）
                accumulator.purchaseFrequency = 30.0 / accumulator.averageInterval;
            } else {
                accumulator.averageInterval = Double.MAX_VALUE;
                accumulator.purchaseFrequency = 0;
            }

            return accumulator;
        }

        @Override
        public PurchaseStatistics merge(PurchaseStatistics a, PurchaseStatistics b) {
            PurchaseStatistics merged = new PurchaseStatistics();
            merged.purchaseCount = a.purchaseCount + b.purchaseCount;
            merged.totalAmount = a.totalAmount + b.totalAmount;
            merged.purchaseTimes.addAll(a.purchaseTimes);
            merged.purchaseTimes.addAll(b.purchaseTimes);
            return merged;
        }
    }

    // 购买统计信息
    public static class PurchaseStatistics {
        int purchaseCount = 0;
        double totalAmount = 0;
        double averageAmount = 0;
        double averageInterval = 0;
        double purchaseFrequency = 0;
        List<Long> purchaseTimes = new ArrayList<>();
    }

    // 生成复购指标
    public static class RepurchaseIndicatorGenerator extends ProcessWindowFunction<PurchaseStatistics, RepurchaseIndicator, String, TimeWindow> {
        @Override
        public void process(String userId, Context context, Iterable<PurchaseStatistics> elements, Collector<RepurchaseIndicator> out) throws Exception {
            PurchaseStatistics stats = elements.iterator().next();
            
            // 计算复购分数（综合指标）
            double frequencyScore = Math.min(stats.purchaseFrequency / 5.0, 1.0); // 最高5次/月
            double amountScore = Math.min(stats.averageAmount / 500.0, 1.0); // 最高500元/单
            double repurchaseScore = 0.6 * frequencyScore + 0.4 * amountScore; // 加权平均
            
            RepurchaseIndicator indicator = new RepurchaseIndicator(
                userId,
                stats.purchaseCount,
                stats.totalAmount,
                stats.averageAmount,
                stats.purchaseFrequency,
                repurchaseScore
            );
            
            out.collect(indicator);
        }
    }
}